#include <sstream>
#include <thread>
#include <fstream>
#include <numeric>
#include "evdeploy/algo/license_plate_rec.h"
#include "evdeploy/utils/report/tool_report.h"
#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>

namespace ev
{
    namespace algo
    {
        std::mutex LicensePlateRec::m_mutex;

        PLATE_COLOR_TYPE ProcessColorType(float *pdata, float &type_score)
        {
            float sum = 0.f;
            int max_index = 0;
            for (int i = 0; i < 8; ++i)
            {
                sum += exp(pdata[i]);
                if (exp(pdata[i]) > exp(pdata[max_index]))
                {
                    max_index = i;
                }
            }
            if (exp(pdata[max_index]) < MINIMUM_POSITIVE_FLOAT)
            {
                type_score = 0;
            }
            else
            {
                type_score = exp(pdata[max_index]) / sum;
            }
            switch (max_index)
            {
            case 0:
                return PLATE_COLOR_TYPE::WHITE_P;
                break;
            case 1:
                return PLATE_COLOR_TYPE::GREEN_P;
                break;
            case 2:
                return PLATE_COLOR_TYPE::YELLOW_P;
                break;
            case 3:
                return PLATE_COLOR_TYPE::BLUE_P;
                break;
            case 4:
                return PLATE_COLOR_TYPE::BLACK_P;
                break;
            case 5:
                return PLATE_COLOR_TYPE::YELGRE_P;
                break;
            case 6:
                return PLATE_COLOR_TYPE::GRAY;
                break;
            case 7:
                return PLATE_COLOR_TYPE::BLUR;
                break;
            default:
                LOG(ERROR) << "error color type";
                return PLATE_COLOR_TYPE::BLUE_P;
                break;
            }
        }

        PLATE_LAYER_TYPE ProcessLayerType(float *pdata, float &type_score)
        {
            float sum = 0.f;
            int max_index = 0;
            for (int i = 0; i < 3; ++i)
            {
                sum += exp(pdata[i]);
                if (exp(pdata[i]) > exp(pdata[max_index]))
                {
                    max_index = i;
                }
            }
            if (exp(pdata[max_index]) < MINIMUM_POSITIVE_FLOAT)
            {
                type_score = 0;
            }
            else
            {
                type_score = exp(pdata[max_index]) / sum;
            }
            switch (max_index)
            {
            case 0:
                return PLATE_LAYER_TYPE::DOUBLE_LAYER;
                break;
            case 1:
                return PLATE_LAYER_TYPE::SINGLE_LAYER;
                break;
            case 2:
                return PLATE_LAYER_TYPE::BLUR;
                break;
            default:
                LOG(ERROR) << "error layer type";
                return PLATE_LAYER_TYPE::SINGLE_LAYER;
                break;
            }
        }

        float ProcessLicense(float *pdata, std::string &license, int &iNumChars)
        {
            license = "";
            std::vector<float> ps{};
            bool bAcceptRepeat = false;

            std::string strLastChar = "-";
            for (int i = 0; i < 24; ++i)
            {
                float tscore = 0;
                int start_index = i;
                int best_index = 0;
                for (int j = 0; j < 76; ++j)
                {
                    tscore += exp(pdata[start_index + j * 24]);
                    if (pdata[start_index + j * 24] > pdata[start_index + best_index * 24])
                    {
                        best_index = j;
                    }
                }
                if (CH_PLATE_CODE[best_index] == "-")
                {
                    bAcceptRepeat = true;
                }
                else
                {
                    if (bAcceptRepeat == true)
                    {
                        license += CH_PLATE_CODE[best_index];
                        strLastChar = CH_PLATE_CODE[best_index];
                        ps.push_back(exp(pdata[start_index + best_index * 24]) / tscore);
                    }
                    else
                    {
                        if (license.size() == 0 || CH_PLATE_CODE[best_index] != strLastChar)
                        {
                            license += CH_PLATE_CODE[best_index];
                            strLastChar = CH_PLATE_CODE[best_index];
                            ps.push_back(exp(pdata[start_index + best_index * 24]) / tscore);
                        }
                    }
                    bAcceptRepeat = false;
                }
            }
            float sum = std::accumulate(ps.begin(), ps.end(), 0.);
            iNumChars = ps.size();
            return ps.size() > 0 ? sum : 0;
        }

        cv::Mat Preprocess(cv::Mat &in_mat)
        {
            cv::Mat resized_mat;
            cv::resize(in_mat, resized_mat, cv::Size(100, 40)); // model input resolution 100x40

            // resized_mat.convertTo(resized_mat, CV_32F, 1 / 255.);
            resized_mat.convertTo(resized_mat, CV_32FC3, 1 / 255.);

            cv::Mat cv_in_mat2 = resized_mat.clone();
            std::vector<cv::Mat> chw_wrappers;
            chw_wrappers.emplace_back(40, 100, CV_32FC1, cv_in_mat2.data);
            chw_wrappers.emplace_back(40, 100, CV_32FC1, cv_in_mat2.data + sizeof(float) * 100 * 40);
            chw_wrappers.emplace_back(40, 100, CV_32FC1, cv_in_mat2.data + 2 * sizeof(float) * 100 * 40);
            cv::split(resized_mat, chw_wrappers);

            return cv_in_mat2;
        }

        LicensePlateRec::LicensePlateRec(std::string engine, std::string uuid)
        {
            if (engine == "atlas" || engine == "trt" || engine == "openvino")
            {
                m_engine = engine;
            }
            else
            {
                EVLOG(ERROR) << "engine not support:" << engine << " , please use:1. atlas, 2. trt, 3. openvino";
                exit(-1);
            }
            m_uuid = uuid;
            EVLOG(INFO) << "use engine:" << m_engine << " and uuid:" << m_uuid;

            ToolReport report;
            // report info
            int code = report.report(2,
                                     "ev_base_model",
                                     "V1.0.0",
                                     "license plate recognition");
            if (code == 0)
            {
                EVLOG(INFO) << "report done successully";
            }
        }

        LicensePlateRec::~LicensePlateRec()
        {
        }

        EVStatus LicensePlateRec::Run(cv::Mat &input_mat, ev::Value &out_json)
        {
            REC_TIME(t0);
            cv::Mat cv_in_mat1 = Preprocess(input_mat);
            EVModelData in;
            EVModelData out;
            EVMatData in_mat;
            in.desc = NULL;
            in.mat = &in_mat;
            in.mat_num = 1;
            in_mat.data = cv_in_mat1.data;
            in_mat.data_size = cv_in_mat1.cols * cv_in_mat1.rows * 3 * 4;
            in_mat.width = cv_in_mat1.cols;
            in_mat.height = cv_in_mat1.rows;
            in_mat.aligned_width = cv_in_mat1.cols;
            in_mat.aligned_height = cv_in_mat1.rows;
            in_mat.channel = 3;
            in_mat.type = EV_UINT8;
            std::lock_guard<std::mutex> lock_guard(m_mutex);
            EVDeploy::GetModel().RunInfer(m_uuid, &in, &out);

            for (int i = 0; i < out.mat_num; ++i)
            {
                std::string dims = "";
                for (int j = 0; j < out.mat[i].dims.size(); ++j)
                {
                    dims += std::to_string(out.mat[i].dims[j]) + ",";
                }
                EVLOG(INFO) << "out " << i << ":" << dims;
            }

            // note output order problem
            float fLicColorTypeScore;
            PLATE_COLOR_TYPE ctype;
            if ((m_engine == "atlas") || (m_engine == "openvino"))
                ctype = ProcessColorType((float *)out.mat[0].data, fLicColorTypeScore);
            else if (m_engine == "trt")
                ctype = ProcessColorType((float *)out.mat[3].data, fLicColorTypeScore);
            //
            float fLicLayerTypeScore;
            PLATE_LAYER_TYPE ltype;
            if ((m_engine == "atlas") || (m_engine == "openvino"))
                ltype = ProcessLayerType((float *)out.mat[1].data, fLicLayerTypeScore);
            else if (m_engine == "trt")
                ltype = ProcessLayerType((float *)out.mat[2].data, fLicLayerTypeScore);

            //
            std::string strUpPartLic;
            int iUpCharNums;
            float fUpPartScore;
            if ((m_engine == "atlas") || (m_engine == "openvino"))
                fUpPartScore = ProcessLicense((float *)out.mat[2].data, strUpPartLic, iUpCharNums);
            else if (m_engine == "trt")
                fUpPartScore = ProcessLicense((float *)out.mat[0].data, strUpPartLic, iUpCharNums);

            //
            std::string strDownPartLic;
            int iDownCharNums;
            float fDownPartScore;
            if ((m_engine == "atlas") || (m_engine == "openvino"))
                fDownPartScore = ProcessLicense((float *)out.mat[3].data, strDownPartLic, iDownCharNums);
            else if (m_engine == "trt")
                fDownPartScore = ProcessLicense((float *)out.mat[1].data, strDownPartLic, iDownCharNums);

            if (iUpCharNums + iDownCharNums > 0)
            {
                out_json["plate_ocr"] = strUpPartLic + strDownPartLic;
                out_json["plate_ocr_score"] = (fUpPartScore + fDownPartScore) / (iUpCharNums + iDownCharNums);
            }
            else
            {
                out_json["plate_ocr"] = "";
                out_json["plate_ocr_score"] = 0;
            }
            out_json["plate_color"] = COLOR_NAME[ctype];
            out_json["plate_color_score"] = fLicColorTypeScore;

            out_json["plate_layer"] = LAYER_NAME[ltype];
            out_json["plate_layer_score"] = fLicLayerTypeScore;

            std::string ret_json = out_json.toStyledString();

            EVLOG(INFO) << "plate rec:\n"
                        << ret_json;

            // 注意释放out.mat,否则会有内存泄露!!!!
            if (out.mat)
            {

                delete[] out.mat;
            }
            REC_TIME(t1);
            EVLOG(INFO) << "license plate rec run time(ms):" << RUN_TIME(t1 - t0);
            return EV_SUCCESS;
        }

    } // namespace algo
} // namespace ev
